# 安装必要的包（取消注释运行）
# pip install llama-index-llms-huggingface
## 查看 gpu运行情况  安装后 控制台运行 nvitop
# pip install nvitop
import os

from llama_index.core.base.llms.types import ChatMessage
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from modelscope import snapshot_download
from llama_index.llms.huggingface import HuggingFaceLLM

# 设置 ModelScope 模型缓存路径
# 默认路径是 c:\Users\Administrator\.cache\modelscope
# 这里设置为 D:\models，避免占用C盘空间
os.environ['MODELSCOPE_CACHE'] = r"D:\models"

### 下载模型到本地缓存目录
# snapshot_download 函数会从 ModelScope 下载模型文件
# 如果模型已经下载过，会直接返回本地路径，不会重复下载
# 模型大小有几个GB，下载需要一定时间
llm_model_dir = snapshot_download(
    "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"  # 模型在ModelScope上的唯一标识
)
# 打印模型本地路径，方便确认
print(f"llm模型已下载到: {llm_model_dir}")


### 下载向量模型
embedding_model_dir = snapshot_download("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
# 打印模型本地路径，方便确认
print(f"embedding模型已下载到: {embedding_model_dir}")

###初始化 HuggingFace 语言模型
# HuggingFaceLLM 是 LlamaIndex 提供的包装类，用于简化模型使用
llm = HuggingFaceLLM(
    # 模型路径：使用下载后的本地目录路径
    model_name=llm_model_dir,

    # tokenizer路径：通常与模型路径相同
    tokenizer_name=llm_model_dir,

    # 上下文窗口大小：模型一次能处理的最大token数量
    # 2048表示模型最多能处理2048个token的输入
    context_window=2048,

    # 最大生成token数：模型回答时最多生成256个token
    max_new_tokens=256,

    # 模型加载参数
    model_kwargs={
        "trust_remote_code": True,  # 信任远程代码（某些模型需要）
        "torch_dtype": "auto"  # 自动选择数据类型（float16/float32）
    },

    # tokenizer加载参数
    tokenizer_kwargs={
        "trust_remote_code": True  # 信任远程代码
    },

    # 文本生成参数
    generate_kwargs={
        "temperature": 0.7,  # 温度参数：控制随机性（0-1，越高越随机）
        "do_sample": True  # 启用采样：使生成结果更多样化
    }
)

### 初始化向量模型
embedding_model = HuggingFaceEmbedding(
    model_name=embedding_model_dir
)

## llm.complete() - 补全模式 只能单个输入 无记忆 每次都独立
response = llm.complete("请用中文介绍一下人工智能的发展历史。")
print("模型回答:", response)
##llm.chat() - 聊天模式 可以处理对话历史，保持上下文 ，支持系统消息、用户消息、助手消息，可以维护连贯性
#response = llm.chat(messages=[ChatMessage(content="请用中文介绍一下人工智能的发展历史。")])
#print("模型回答:", response)
